import os
from operator import itemgetter

from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.tools import QuerySQLDataBaseTool
from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI

# os.environ['http_proxy'] = '127.0.0.1:7890'
# os.environ['https_proxy'] = '127.0.0.1:7890'

# 聊天机器人案例
# 创建模型
model = ChatOpenAI(model='gpt-3.5-turbo')

# sqlalchemy 初始化MySQL数据库的连接
HOSTNAME = '10.219.8.30'
PORT = '3306'
DATABASE = 'test_go'
USERNAME = 'dev'
PASSWORD = 'dev'
# mysqlclient驱动URL 改为 pymysql的
MYSQL_URI = 'mysql+pymysql://{}:{}@{}:{}/{}?charset=utf8mb4'.format(USERNAME, PASSWORD, HOSTNAME, PORT, DATABASE)

db = SQLDatabase.from_uri(MYSQL_URI)

# 获取可用的表名
usable_tables = db.get_usable_table_names()
print(f"Usable tables: {usable_tables}")

# 选择与查询相关的表
relevant_tables = ['app_feedback']  # 这里只选择需要的表
# 打印 relevant_tables 以确认其内容
print(f"Relevant tables: {relevant_tables}")

# 精简数据库模式信息
# table_info = {table: db.get_table_info('app_feedback') for table in relevant_tables}
# table_info = relevant_tables

# 测试连接是否成功
# print(db.get_usable_table_names())
# print(db.run('select * from app_feedback limit 10;'))

# 直接使用大模型和数据库整合, 只能根据你的问题生成SQL
# 初始化生成SQL的chain
test_chain = create_sql_query_chain(model, db)
# 提问
# resp = test_chain.invoke({'question': '请问：问题反馈表中有多少条数据？'})
# 打印生成的SQL
# print(resp)


answer_prompt = PromptTemplate.from_template("""给定以下用户问题、SQL语句和SQL执行后的结果，回答用户问题。
    Question: {question}
    SQL Query: {query}
    SQL Result: {result}
    回答: """)
# # 创建一个执行sql语句的工具
execute_sql_tool = QuerySQLDataBaseTool(db=db)
#
# # 1、生成SQL，2、执行SQL
# # 2、模板
chain = (RunnablePassthrough
         .assign(query=test_chain)
         .assign(result=itemgetter('query') | execute_sql_tool)
         | answer_prompt | model | StrOutputParser())

rep = chain.invoke(input={'question': '请问：问题反馈表中有多少条数据？'})
print(rep)